# -*- coding: utf-8 -*-
import numpy as np

# 卡尔曼滤波预测类，用于预测 UAV 的位置和速度
class KalmanFilterPredictor:
    # 处理值噪声越大，越依赖观测
    def __init__(self, dt=0.1, process_noise=1, measurement_noise=0.1):
        self.x = None  # 初始状态（位置和速度）

        # 状态转移矩阵 初始化的dt不需要被用到
        self.F = np.array([[1, 0, 0, dt, 0, 0],
                           [0, 1, 0, 0, dt, 0],
                           [0, 0, 1, 0, 0, dt],
                           [0, 0, 0, 1, 0, 0],
                           [0, 0, 0, 0, 1, 0],
                           [0, 0, 0, 0, 0, 1]], dtype=float)

        # 测量矩阵
        self.H = np.array([[1, 0, 0, 0, 0, 0],
                           [0, 1, 0, 0, 0, 0],
                           [0, 0, 1, 0, 0, 0]], dtype=float)

        # 初始状态协方差矩阵
        self.P = np.eye(6)
        # 测量噪声协方差矩阵
        self.R = np.eye(3) * measurement_noise 
        # 过程误差协方差矩阵
        process_noise_vel = 0.4
        self.Q = np.diag([process_noise, process_noise, process_noise, process_noise_vel, process_noise_vel, process_noise_vel])
        # self.Q = np.eye(6) * 0.1 
        
        # # 初始状态协方差矩阵
        # self.P = np.eye(6)
        # # 测量噪声协方差矩阵
        # self.R = np.eye(3)
        # # 过程误差协方差矩阵
        # self.Q = np.eye(6) * 0.1

        # 存储滤波后的位置和速度
        self.filtered_positions = []
        self.filtered_velocities = []

    def initialize_state(self, initial_position):
        # 初始化状态向量
        self.x = np.array([initial_position[0], initial_position[1], initial_position[2], 0, 0, 0], dtype=float)
    
    # def initialize_state(self, initial_position,  initail_vel):
    #     # 初始化状态向量
    #     self.x = np.array([initial_position[0], initial_position[1], initial_position[2], initail_vel[0], initail_vel[1], initail_vel[2]], dtype=float)

    def predict(self, z, dt):
        # 更新时间间隔
        self.F[0, 3] = dt
        self.F[1, 4] = dt
        self.F[2, 5] = dt
        
        # 预测步骤
        self.x = self.F @ self.x
        print('x_prior:', self.x)
        self.P = self.F @ self.P @ self.F.T + self.Q
        print('P_prior:',self.P)

        # 测量更新步骤
        y = np.array(z) - self.H @ self.x
        print('y:', y)
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        print('K:',K)
        self.x = self.x + K @ y
        print('x_esti:', self.x)
        self.P = (np.eye(6) - K @ self.H) @ self.P
        
        print('P_esti:',self.P)
        print('R:',self.R)
        print('Q:',self.Q)

        # 存储结果
        self.filtered_positions.clear()
        self.filtered_velocities.clear()
        self.filtered_positions.append((self.x[0], self.x[1], self.x[2]))
        self.filtered_velocities.append((self.x[3], self.x[4], self.x[5]))
